{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# How to Debug Networks\n",
    "\n",
    "This guide shows you how to identify and fix common issues when developing neural networks with BrainPy."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Quick Diagnostic Checklist\n",
    "\n",
    "When your network isn't working, check these first:\n",
    "\n",
    "**☐ Is the network receiving input?**  \n",
    "   Print input values, check shapes\n",
    "\n",
    "**☐ Are neurons firing?**  \n",
    "   Count spikes, check spike rates\n",
    "\n",
    "**☐ Are projections working?**  \n",
    "   Verify connectivity, check weights\n",
    "\n",
    "**☐ Is update order correct?**  \n",
    "   Get spikes BEFORE updating neurons\n",
    "\n",
    "**☐ Are states initialized?**  \n",
    "   Call `brainstate.nn.init_all_states()`\n",
    "\n",
    "**☐ Are units correct?**  \n",
    "   All values need physical units (mV, nA, ms)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Common Issues and Solutions\n",
    "\n",
    "### Issue 1: No Spikes / Silent Network\n",
    "\n",
    "**Symptoms:**\n",
    "\n",
    "- Network produces no spikes\n",
    "- All neurons stay at rest potential\n",
    "\n",
    "**Diagnosis:**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import brainpy\n",
    "import brainstate\n",
    "import brainunit as u\n",
    "import jax.numpy as jnp\n",
    "\n",
    "neuron = brainpy.state.LIF(100, V_rest=-65*u.mV, V_th=-50*u.mV, tau=10*u.ms)\n",
    "brainstate.nn.init_all_states(neuron)\n",
    "\n",
    "# Check 1: Is input being provided?\n",
    "inp = brainstate.random.rand(100) * 5.0 * u.nA\n",
    "print(\"Input range:\", inp.min(), \"to\", inp.max())\n",
    "\n",
    "# Check 2: Are neurons updating?\n",
    "V_before = neuron.V.value.copy()\n",
    "neuron(inp)\n",
    "V_after = neuron.V.value\n",
    "print(\"Voltage changed:\", not jnp.allclose(V_before, V_after))\n",
    "\n",
    "# Check 3: Are any neurons near threshold?\n",
    "print(\"Max voltage:\", V_after.max())\n",
    "print(\"Threshold:\", neuron.V_th.to_decimal(u.mV))\n",
    "print(\"Neurons above -55mV:\", jnp.sum(V_after > -55))\n",
    "\n",
    "# Check 4: Count spikes\n",
    "for i in range(100):\n",
    "    neuron(inp)\n",
    "spike_count = jnp.sum(neuron.spike.value)\n",
    "print(f\"Spikes in 100 steps: {spike_count}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Common Causes:**\n",
    "\n",
    "1. **Input too weak:**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Too weak\n",
    "inp = brainstate.random.rand(100) * 0.1 * u.nA  # Not enough!\n",
    "\n",
    "# Better\n",
    "inp = brainstate.random.rand(100) * 5.0 * u.nA  # Stronger"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "2. **Threshold too high:**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Check threshold\n",
    "neuron = brainpy.state.LIF(100, V_th=-40*u.mV, ...)  # Harder to spike\n",
    "neuron = brainpy.state.LIF(100, V_th=-50*u.mV, ...)  # Easier to spike"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "3. **Time constant too large:**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Slow integration\n",
    "neuron = brainpy.state.LIF(100, tau=100*u.ms, ...)  # Very slow\n",
    "\n",
    "# Faster\n",
    "neuron = brainpy.state.LIF(100, tau=10*u.ms, ...)  # Normal speed"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "4. **Missing initialization:**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "neuron = brainpy.state.LIF(100, ...)\n",
    "# MUST initialize!\n",
    "brainstate.nn.init_all_states(neuron)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Issue 2: Runaway Activity / Explosion\n",
    "\n",
    "**Symptoms:**\n",
    "\n",
    "- All neurons fire constantly\n",
    "- Membrane potentials go to infinity\n",
    "- NaN values appear\n",
    "\n",
    "**Diagnosis:**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# Check for NaN\n",
    "if jnp.any(jnp.isnan(neuron.V.value)):\n",
    "    print(\"❌ NaN detected in membrane potential!\")\n",
    "\n",
    "# Check for explosion\n",
    "if jnp.any(jnp.abs(neuron.V.value) > 1000):\n",
    "    print(\"❌ Membrane potential exploded!\")\n",
    "\n",
    "# Check spike rate\n",
    "spike_rate = jnp.mean(neuron.spike.value)\n",
    "print(f\"Spike rate: {spike_rate*100:.1f}%\")\n",
    "if spike_rate > 0.5:\n",
    "    print(\"⚠️ More than 50% of neurons firing every step!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Common Causes:**\n",
    "\n",
    "1. **Excitation-Inhibition imbalance:**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Imbalanced (explosion!)\n",
    "w_exc = 5.0 * u.mS  # Too strong\n",
    "w_inh = 1.0 * u.mS  # Too weak\n",
    "\n",
    "# Balanced\n",
    "w_exc = 0.5 * u.mS\n",
    "w_inh = 5.0 * u.mS  # Inhibition ~10× stronger"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "2. **Positive feedback loop:**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Check recurrent excitation\n",
    "# E → E with no inhibition can explode\n",
    "\n",
    "# Add inhibition\n",
    "class BalancedNetwork(brainstate.nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.E = brainpy.state.LIF(800, ...)\n",
    "        self.I = brainpy.state.LIF(200, ...)\n",
    "\n",
    "        self.E2E = ...  # Excitatory recurrence\n",
    "        self.I2E = ...  # MUST have inhibition!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "3. **Time step too large:**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Unstable\n",
    "brainstate.environ.set(dt=1.0 * u.ms)  # Too large\n",
    "\n",
    "# Stable\n",
    "brainstate.environ.set(dt=0.1 * u.ms)  # Standard"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "4. **Wrong reversal potentials:**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# WRONG: Inhibition with excitatory reversal\n",
    "out_inh = brainpy.state.COBA.desc(E=0*u.mV)  # Should be negative!\n",
    "\n",
    "# CORRECT\n",
    "out_exc = brainpy.state.COBA.desc(E=0*u.mV)    # Excitation\n",
    "out_inh = brainpy.state.COBA.desc(E=-80*u.mV)  # Inhibition"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Issue 3: Spikes Not Propagating\n",
    "\n",
    "**Symptoms:**\n",
    "\n",
    "- Presynaptic neurons spike\n",
    "- Postsynaptic neurons don't respond\n",
    "- Projection seems inactive\n",
    "\n",
    "**Diagnosis:**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create simple network\n",
    "pre = brainpy.state.LIF(10, V_rest=-65*u.mV, V_th=-50*u.mV, tau=10*u.ms)\n",
    "post = brainpy.state.LIF(10, V_rest=-65*u.mV, V_th=-50*u.mV, tau=10*u.ms)\n",
    "\n",
    "proj = brainpy.state.AlignPostProj(\n",
    "    comm=brainstate.nn.EventFixedProb(10, 10, prob=0.5, weight=2.0*u.mS),\n",
    "    syn=brainpy.state.Expon.desc(10, tau=5*u.ms),\n",
    "    out=brainpy.state.CUBA.desc(),\n",
    "    post=post\n",
    ")\n",
    "\n",
    "brainstate.nn.init_all_states([pre, post, proj])\n",
    "\n",
    "# Diagnosis\n",
    "for i in range(10):\n",
    "    # CRITICAL: Get spikes BEFORE update\n",
    "    pre_spikes = pre.get_spike()\n",
    "\n",
    "    # Strong input to pre\n",
    "    pre(brainstate.random.rand(10) * 10.0 * u.nA)\n",
    "\n",
    "    # Check: Did pre spike?\n",
    "    if jnp.sum(pre_spikes) > 0:\n",
    "        print(f\"Step {i}: {jnp.sum(pre_spikes)} presynaptic spikes\")\n",
    "\n",
    "        # Update projection\n",
    "        proj(pre_spikes)\n",
    "\n",
    "        # Check: Did projection produce current?\n",
    "        print(f\"  Synaptic conductance: {proj.syn.g.value.max():.4f}\")\n",
    "\n",
    "    # Update post\n",
    "    post(0*u.nA)  # Only synaptic input\n",
    "\n",
    "    # Check: Did post spike?\n",
    "    post_spikes = post.get_spike()\n",
    "    print(f\"  {jnp.sum(post_spikes)} postsynaptic spikes\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Common Causes:**\n",
    "\n",
    "1. **Wrong spike timing:**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# WRONG: Spikes from current step\n",
    "pre(inp)              # Update first\n",
    "spikes = pre.get_spike()  # These are NEW spikes\n",
    "proj(spikes)          # But projection needs OLD spikes!\n",
    "\n",
    "# CORRECT: Spikes from previous step\n",
    "spikes = pre.get_spike()  # Get OLD spikes first\n",
    "proj(spikes)              # Update projection\n",
    "pre(inp)                  # Then update neurons"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "2. **Weak connectivity:**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Too sparse\n",
    "comm = brainstate.nn.EventFixedProb(..., prob=0.01, weight=0.1*u.mS)\n",
    "\n",
    "# Stronger\n",
    "comm = brainstate.nn.EventFixedProb(..., prob=0.1, weight=1.0*u.mS)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "3. **Missing projection update:**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Forgot to call projection!\n",
    "spk = pre.get_spike()\n",
    "# proj(spk)  <- MISSING!\n",
    "post(0*u.nA)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "4. **Wrong postsynaptic target:**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Wrong target\n",
    "proj = brainpy.state.AlignPostProj(..., post=wrong_population)\n",
    "\n",
    "# Correct target\n",
    "proj = brainpy.state.AlignPostProj(..., post=correct_population)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Issue 4: Shape Mismatch Errors\n",
    "\n",
    "**Symptoms:**\n",
    "\n",
    "```\n",
    "ValueError: operands could not be broadcast together\n",
    "with shapes (100,) (64, 100)\n",
    "```\n",
    "\n",
    "**Common Causes:**\n",
    "\n",
    "1. **Batch dimension mismatch:**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Network initialized with batch\n",
    "brainstate.nn.init_all_states(net, batch_size=64)\n",
    "# States shape: (64, 100)\n",
    "\n",
    "# But input has no batch\n",
    "inp = jnp.zeros(100)  # Shape: (100,) - WRONG!\n",
    "\n",
    "# Fix: Add batch dimension\n",
    "inp = jnp.zeros((64, 100))  # Shape: (64, 100) - CORRECT"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "2. **Forgot batch in initialization:**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialized without batch\n",
    "brainstate.nn.init_all_states(net)  # Shape: (100,)\n",
    "\n",
    "# But providing batched input\n",
    "inp = jnp.zeros((64, 100))  # Shape: (64, 100)\n",
    "\n",
    "# Fix: Initialize with batch\n",
    "brainstate.nn.init_all_states(net, batch_size=64)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Debug shape mismatches:**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"Input shape: {inp.shape}\")\n",
    "print(f\"Network state shape: {net.neurons.V.value.shape}\")\n",
    "print(f\"Expected: Both should have same batch dimension\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Inspection Tools\n",
    "\n",
    "### Print State Values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Inspect neuron states\n",
    "neuron = brainpy.state.LIF(10, ...)\n",
    "brainstate.nn.init_all_states(neuron)\n",
    "\n",
    "print(\"Membrane potentials:\", neuron.V.value)\n",
    "print(\"Spikes:\", neuron.spike.value)\n",
    "print(\"Shape:\", neuron.V.value.shape)\n",
    "\n",
    "# Statistics\n",
    "print(f\"V range: [{neuron.V.value.min():.2f}, {neuron.V.value.max():.2f}]\")\n",
    "print(f\"V mean: {neuron.V.value.mean():.2f}\")\n",
    "print(f\"Spike count: {jnp.sum(neuron.spike.value)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Visualize Activity"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "# Record activity\n",
    "n_steps = 1000\n",
    "V_history = []\n",
    "spike_history = []\n",
    "\n",
    "for i in range(n_steps):\n",
    "    neuron(inp)\n",
    "    V_history.append(neuron.V.value.copy())\n",
    "    spike_history.append(neuron.spike.value.copy())\n",
    "\n",
    "V_history = jnp.array(V_history)\n",
    "spike_history = jnp.array(spike_history)\n",
    "\n",
    "# Plot membrane potential\n",
    "plt.figure(figsize=(12, 4))\n",
    "plt.plot(V_history[:, 0])  # First neuron\n",
    "plt.xlabel('Time step')\n",
    "plt.ylabel('Membrane Potential (mV)')\n",
    "plt.title('Neuron 0 Membrane Potential')\n",
    "plt.show()\n",
    "\n",
    "# Plot raster\n",
    "plt.figure(figsize=(12, 6))\n",
    "times, neurons = jnp.where(spike_history > 0)\n",
    "plt.scatter(times, neurons, s=1, c='black')\n",
    "plt.xlabel('Time step')\n",
    "plt.ylabel('Neuron index')\n",
    "plt.title('Spike Raster')\n",
    "plt.show()\n",
    "\n",
    "# Firing rate over time\n",
    "plt.figure(figsize=(12, 4))\n",
    "firing_rate = jnp.mean(spike_history, axis=1) * 1000 / 0.1  # Hz\n",
    "plt.plot(firing_rate)\n",
    "plt.xlabel('Time step')\n",
    "plt.ylabel('Population Rate (Hz)')\n",
    "plt.title('Population Firing Rate')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Check Connectivity"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# For sparse projections\n",
    "proj = brainpy.state.AlignPostProj(\n",
    "    comm=brainstate.nn.EventFixedProb(100, 50, prob=0.1, weight=0.5*u.mS),\n",
    "    syn=brainpy.state.Expon.desc(50, tau=5*u.ms),\n",
    "    out=brainpy.state.CUBA.desc(),\n",
    "    post=post_neurons\n",
    ")\n",
    "\n",
    "# Check connection count\n",
    "print(f\"Expected connections: {100 * 50 * 0.1:.0f}\")\n",
    "# Note: Actual connectivity may vary due to randomness\n",
    "\n",
    "# Check weights\n",
    "# (Accessing internal connectivity structure depends on implementation)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Monitor Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Track loss and metrics\n",
    "train_losses = []\n",
    "val_accuracies = []\n",
    "\n",
    "for epoch in range(num_epochs):\n",
    "    epoch_losses = []\n",
    "\n",
    "    for batch in train_loader:\n",
    "        loss = train_step(net, batch)\n",
    "        epoch_losses.append(float(loss))\n",
    "\n",
    "    avg_loss = np.mean(epoch_losses)\n",
    "    train_losses.append(avg_loss)\n",
    "\n",
    "    # Validation\n",
    "    val_acc = evaluate(net, val_loader)\n",
    "    val_accuracies.append(val_acc)\n",
    "\n",
    "    print(f\"Epoch {epoch}: Loss={avg_loss:.4f}, Val Acc={val_acc:.2%}\")\n",
    "\n",
    "    # Check for issues\n",
    "    if np.isnan(avg_loss):\n",
    "        print(\"❌ NaN loss! Stopping training.\")\n",
    "        break\n",
    "\n",
    "    if avg_loss > 10 * train_losses[0]:\n",
    "        print(\"⚠️ Loss exploding!\")\n",
    "\n",
    "# Plot training curves\n",
    "plt.figure(figsize=(12, 4))\n",
    "plt.subplot(1, 2, 1)\n",
    "plt.plot(train_losses)\n",
    "plt.xlabel('Epoch')\n",
    "plt.ylabel('Loss')\n",
    "plt.title('Training Loss')\n",
    "\n",
    "plt.subplot(1, 2, 2)\n",
    "plt.plot(val_accuracies)\n",
    "plt.xlabel('Epoch')\n",
    "plt.ylabel('Accuracy')\n",
    "plt.title('Validation Accuracy')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Advanced Debugging\n",
    "\n",
    "### Gradient Checking"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import braintools\n",
    "\n",
    "# Check if gradients are being computed\n",
    "params = net.states(brainstate.ParamState)\n",
    "\n",
    "grads = brainstate.transform.grad(\n",
    "    loss_fn,\n",
    "    params,\n",
    "    return_value=True\n",
    ")(net, X, y)\n",
    "\n",
    "# Inspect gradients\n",
    "for name, grad in grads.items():\n",
    "    grad_norm = jnp.linalg.norm(grad.value.flatten())\n",
    "    print(f\"{name}: gradient norm = {grad_norm:.6f}\")\n",
    "\n",
    "    if jnp.any(jnp.isnan(grad.value)):\n",
    "        print(f\"  ❌ NaN in gradient!\")\n",
    "\n",
    "    if grad_norm == 0:\n",
    "        print(f\"  ⚠️ Zero gradient - parameter not learning\")\n",
    "\n",
    "    if grad_norm > 1000:\n",
    "        print(f\"  ⚠️ Exploding gradient!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Trace Execution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def debug_step(net, inp):\n",
    "    \"\"\"Instrumented simulation step.\"\"\"\n",
    "    print(f\"\\n--- Step Start ---\")\n",
    "\n",
    "    # Before\n",
    "    print(f\"Input range: [{inp.min():.2f}, {inp.max():.2f}]\")\n",
    "    print(f\"V before: [{net.neurons.V.value.min():.2f}, {net.neurons.V.value.max():.2f}]\")\n",
    "\n",
    "    # Execute\n",
    "    output = net(inp)\n",
    "\n",
    "    # After\n",
    "    print(f\"V after: [{net.neurons.V.value.min():.2f}, {net.neurons.V.value.max():.2f}]\")\n",
    "    print(f\"Spikes: {jnp.sum(net.neurons.spike.value)}\")\n",
    "    print(f\"Output range: [{output.min():.2f}, {output.max():.2f}]\")\n",
    "\n",
    "    # Checks\n",
    "    if jnp.any(jnp.isnan(net.neurons.V.value)):\n",
    "        print(\"❌ NaN detected!\")\n",
    "        import pdb; pdb.set_trace()  # Drop into debugger\n",
    "\n",
    "    print(f\"--- Step End ---\\n\")\n",
    "    return output\n",
    "\n",
    "# Use for debugging\n",
    "for i in range(10):\n",
    "    output = debug_step(net, input_data)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Assertion Checks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SafeNetwork(brainstate.nn.Module):\n",
    "    \"\"\"Network with built-in checks.\"\"\"\n",
    "\n",
    "    def __init__(self, n_neurons=100):\n",
    "        super().__init__()\n",
    "        self.neurons = brainpy.state.LIF(n_neurons, ...)\n",
    "\n",
    "    def update(self, inp):\n",
    "        # Pre-checks\n",
    "        assert inp.shape[-1] == 100, f\"Wrong input size: {inp.shape}\"\n",
    "        assert not jnp.any(jnp.isnan(inp)), \"NaN in input!\"\n",
    "        assert not jnp.any(jnp.isinf(inp)), \"Inf in input!\"\n",
    "\n",
    "        # Execute\n",
    "        self.neurons(inp)\n",
    "        output = self.neurons.get_spike()\n",
    "\n",
    "        # Post-checks\n",
    "        assert not jnp.any(jnp.isnan(self.neurons.V.value)), \"NaN in membrane potential!\"\n",
    "        assert jnp.all(jnp.abs(self.neurons.V.value) < 1000), \"Voltage explosion!\"\n",
    "\n",
    "        return output"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Unit Testing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def test_neuron_spikes():\n",
    "    \"\"\"Test that neuron spikes with strong input.\"\"\"\n",
    "    neuron = brainpy.state.LIF(1, V_rest=-65*u.mV, V_th=-50*u.mV, tau=10*u.ms)\n",
    "    brainstate.nn.init_all_states(neuron)\n",
    "\n",
    "    # Strong constant input should cause spiking\n",
    "    strong_input = jnp.array([20.0]) * u.nA\n",
    "\n",
    "    spike_count = 0\n",
    "    for _ in range(100):\n",
    "        neuron(strong_input)\n",
    "        spike_count += int(neuron.spike.value[0])\n",
    "\n",
    "    assert spike_count > 0, \"Neuron didn't spike with strong input!\"\n",
    "    assert spike_count < 100, \"Neuron spiked every step (check reset!)\"\n",
    "\n",
    "    print(f\"✅ Neuron test passed ({spike_count} spikes)\")\n",
    "\n",
    "def test_projection():\n",
    "    \"\"\"Test that projection propagates spikes.\"\"\"\n",
    "    pre = brainpy.state.LIF(10, V_rest=-65*u.mV, V_th=-50*u.mV, tau=10*u.ms)\n",
    "    post = brainpy.state.LIF(10, V_rest=-65*u.mV, V_th=-50*u.mV, tau=10*u.ms)\n",
    "\n",
    "    proj = brainpy.state.AlignPostProj(\n",
    "        comm=brainstate.nn.EventFixedProb(10, 10, prob=1.0, weight=5.0*u.mS),  # 100% connectivity\n",
    "        syn=brainpy.state.Expon.desc(10, tau=5*u.ms),\n",
    "        out=brainpy.state.CUBA.desc(),\n",
    "        post=post\n",
    "    )\n",
    "\n",
    "    brainstate.nn.init_all_states([pre, post, proj])\n",
    "\n",
    "    # Make pre spike\n",
    "    pre(jnp.ones(10) * 20.0 * u.nA)\n",
    "\n",
    "    # Projection should activate\n",
    "    spk = pre.get_spike()\n",
    "    assert jnp.sum(spk) > 0, \"Pre didn't spike!\"\n",
    "\n",
    "    proj(spk)\n",
    "\n",
    "    # Check synaptic conductance increased\n",
    "    assert proj.syn.g.value.max() > 0, \"Synapse didn't activate!\"\n",
    "\n",
    "    print(\"✅ Projection test passed\")\n",
    "\n",
    "# Run tests\n",
    "test_neuron_spikes()\n",
    "test_projection()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Debugging Checklist\n",
    "\n",
    "When your network doesn't work:\n",
    "\n",
    "**1. Check Initialization**\n",
    "\n",
    "```python\n",
    "☐ Called brainstate.nn.init_all_states()?\n",
    "☐ Correct batch_size parameter?\n",
    "☐ All submodules initialized?\n",
    "```\n",
    "\n",
    "**2. Check Input**\n",
    "\n",
    "```python\n",
    "☐ Input shape matches network?\n",
    "☐ Input has units (nA, mV, etc.)?\n",
    "☐ Input magnitude reasonable?\n",
    "☐ Input not all zeros?\n",
    "```\n",
    "\n",
    "**3. Check Neurons**\n",
    "\n",
    "```python\n",
    "☐ Threshold reasonable (e.g., -50 mV)?\n",
    "☐ Reset potential below threshold?\n",
    "☐ Time constant reasonable (5-20 ms)?\n",
    "☐ Neurons actually spiking?\n",
    "```\n",
    "\n",
    "**4. Check Projections**\n",
    "\n",
    "```python\n",
    "☐ Connectivity probability > 0?\n",
    "☐ Weights reasonable magnitude?\n",
    "☐ Correct update order (spikes before update)?\n",
    "☐ Projection actually called?\n",
    "```\n",
    "\n",
    "**5. Check Balance**\n",
    "\n",
    "```python\n",
    "☐ Inhibition stronger than excitation (~10×)?\n",
    "☐ Reversal potentials correct (E=0, I=-80)?\n",
    "☐ E/I ratio appropriate (4:1)?\n",
    "```\n",
    "\n",
    "**6. Check Training**\n",
    "\n",
    "```python\n",
    "☐ Loss decreasing?\n",
    "☐ Gradients non-zero?\n",
    "☐ No NaN in gradients?\n",
    "☐ Learning rate appropriate?\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Common Error Messages\n",
    "\n",
    "### \"operands could not be broadcast\"\n",
    "\n",
    "**Meaning:** Shape mismatch\n",
    "\n",
    "**Fix:** Check batch dimensions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"Shapes: {x.shape} vs {y.shape}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### \"RESOURCE_EXHAUSTED: Out of memory\"\n",
    "\n",
    "**Meaning:** GPU/CPU memory full\n",
    "\n",
    "**Fix:** Reduce batch size or network size"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Reduce batch\n",
    "brainstate.nn.init_all_states(net, batch_size=16)  # Instead of 64"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### \"Concrete value required\"\n",
    "\n",
    "**Meaning:** JIT can't handle dynamic values\n",
    "\n",
    "**Fix:** Use static shapes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Dynamic (bad for JIT)\n",
    "n = len(data)  # Changes each call\n",
    "\n",
    "# Static (good for JIT)\n",
    "n = 100  # Fixed value"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### \"Invalid device\"\n",
    "\n",
    "**Meaning:** Trying to use unavailable device\n",
    "\n",
    "**Fix:** Check available devices"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import jax\n",
    "print(jax.devices())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Best Practices\n",
    "\n",
    "✅ **Test small first** - Debug with 10 neurons before scaling to 10,000\n",
    "\n",
    "✅ **Visualize early** - Plot activity to see problems immediately\n",
    "\n",
    "✅ **Check incrementally** - Test each component before combining\n",
    "\n",
    "✅ **Use assertions** - Catch problems early with runtime checks\n",
    "\n",
    "✅ **Print liberally** - Add diagnostic prints during development\n",
    "\n",
    "✅ **Keep backups** - Save working versions before major changes\n",
    "\n",
    "✅ **Start simple** - Begin with minimal network, add complexity gradually\n",
    "\n",
    "✅ **Write tests** - Unit test individual components\n",
    "\n",
    "❌ **Don't debug by guessing** - Use systematic diagnosis\n",
    "\n",
    "❌ **Don't skip initialization** - Always call init_all_states\n",
    "\n",
    "❌ **Don't ignore warnings** - They often indicate real problems"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Summary\n",
    "\n",
    "**Debugging workflow:**\n",
    "\n",
    "1. **Identify symptom** (no spikes, explosion, etc.)\n",
    "2. **Isolate component** (neurons, projections, input)\n",
    "3. **Inspect state** (print values, plot activity)\n",
    "4. **Form hypothesis** (what might be wrong?)\n",
    "5. **Test fix** (make one change at a time)\n",
    "6. **Verify** (ensure problem solved)\n",
    "\n",
    "**Quick diagnostic code:**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Comprehensive diagnostic\n",
    "def diagnose_network(net, inp):\n",
    "    print(\"=== Network Diagnostic ===\")\n",
    "\n",
    "    # Input\n",
    "    print(f\"Input shape: {inp.shape}\")\n",
    "    print(f\"Input range: [{inp.min():.2f}, {inp.max():.2f}]\")\n",
    "\n",
    "    # States\n",
    "    if hasattr(net, 'neurons'):\n",
    "        V = net.neurons.V.value\n",
    "        print(f\"Voltage shape: {V.shape}\")\n",
    "        print(f\"Voltage range: [{V.min():.2f}, {V.max():.2f}]\")\n",
    "\n",
    "    # Simulation\n",
    "    output = net(inp)\n",
    "\n",
    "    # Results\n",
    "    if hasattr(net, 'neurons'):\n",
    "        spk_count = jnp.sum(net.neurons.spike.value)\n",
    "        print(f\"Spikes: {spk_count}\")\n",
    "\n",
    "    print(f\"Output shape: {output.shape}\")\n",
    "    print(f\"Output range: [{output.min():.2f}, {output.max():.2f}]\")\n",
    "\n",
    "    # Checks\n",
    "    if jnp.any(jnp.isnan(output)):\n",
    "        print(\"❌ NaN in output!\")\n",
    "    if jnp.all(output == 0):\n",
    "        print(\"⚠️  Output all zeros!\")\n",
    "\n",
    "    print(\"=========================\")\n",
    "    return output"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## See Also\n",
    "\n",
    "- Core Concepts: State Management\n",
    "- Core Concepts: Projections\n",
    "- Performance Optimization"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Ecosystem-py",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
